using Zygote
using BenchmarkTools
using KernelFunctions

x = [3.0]

macro old_check_args(K, param, cond, desc=string(cond))
    quote
        if !($(esc(cond)))
            throw(
                ArgumentError(
                    string(
                        $(string(K)),
                        ": ",
                        $(string(param)),
                        " = ",
                        $(esc(param)),
                        " does not ",
                        "satisfy the constraint ",
                        $(string(desc)),
                        ".",
                    ),
                ),
            )
        end
    end
end


struct OldLinearKernel{Tc<:Real} <: KernelFunctions.SimpleKernel
    c::Vector{Tc}

    function OldLinearKernel(c::Real)
        @old_check_args(LinearKernel, c, c >= zero(c), "c ≥ 0")
        return new{typeof(c)}([c])
    end
end

function f(x)
    k = LinearKernel(;c = x[1])
    sum(k.c)
end
function g(x)
    k = LinearKernel(;c = x[1], check_args=false)
    sum(k.c)
end
function h(x)
    k = OldLinearKernel(x[1])
    sum(k.c)
end

@btime Zygote.gradient($f, $x) # 15.980 μs (150 allocations: 5.89 KiB)
@btime Zygote.gradient($g, $x) #  13.853 μs (142 allocations: 5.72 KiB)
@btime Zygote.gradient($h, $x)  #  4.700 μs (51 allocations: 1.89 KiB)

    function OldLinearKernel(c::Real; check_args=true)
        check_args && @old_check_args(LinearKernel, c, c >= zero(c), "c ≥ 0")
        return new{typeof(c)}([c])
    end

function h(x)
    k = OldLinearKernel(x[1])
    sum(k.c)
end
function i(x)
    k = OldLinearKernel(x[1]; check_args=false)
    sum(k.c)
end

@btime Zygote.gradient($h, $x)  # 6.086 μs (65 allocations: 2.58 KiB)
@btime Zygote.gradient($i, $x) # 10.404 μs (91 allocations: 3.38 KiB)